package com.shujia.flink.kafka;

import com.shujia.flink.core.Car;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.HashMap;

public class Demo3KafkaCars {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        //构建kafka source
        KafkaSource<String> source = KafkaSource.<String>builder()
                //指定broker列表
                .setBootstrapServers("master:9092,node1:9092,node2:9092")
                //指定topic
                .setTopics("cars")
                //消费者组
                .setGroupId("Demo3KafkaCars")
                //指定读取数据的位置：earliest：读取最早的数据, latest: 读取最新的数据
                .setStartingOffsets(OffsetsInitializer.latest())
                //读取数据的格式
                .setValueOnlyDeserializer(new SimpleStringSchema())
                .build();

        //使用 kafka source
        DataStreamSource<String> kafkaDS = env
                .fromSource(source, WatermarkStrategy.noWatermarks(), "Kafka Source");


        //解析数据
        SingleOutputStreamOperator<Car> carsDS = kafkaDS.map(line -> {
            String[] split = line.split(",");
            long ts = Long.parseLong(split[7]) * 1000;
            double speed = Double.parseDouble(split[8]);
            return new Car(split[0], split[1], split[2], split[3], split[4], split[5], split[6], ts, speed);
        });

        //实时统计每个道路总的车流量和总的平均车速
        //安装道路分组
        KeyedStream<Car, String> keyByDS = carsDS.keyBy(car -> car.getRoadId());

        SingleOutputStreamOperator<Result> resultsDS = keyByDS.process(new KeyedProcessFunction<String, Car, Result>() {

            //定义在这个的成员变量，同一个task中的key共享同一个
            //Long flow = 0L;
            //Double sumSpeed = 0.0;

            //使用hashmap保存车流量，道路编号作为key,车流量作为value
            HashMap<String, Long> flowMap = new HashMap<>();
            HashMap<String, Double> sumSpeedMap = new HashMap<>();


            /**
             * processElement:每一条数据执行一次
             * @param car： 一行数据
             * @param ctx : 上下文对象
             * @param out ：用于将数据发送到下游
             */
            @Override
            public void processElement(Car car,
                                       KeyedProcessFunction<String, Car, Result>.Context ctx,
                                       Collector<Result> out) throws Exception {

                String roadId = car.getRoadId();

                //从hashmap中通过道路编号获取车流量
                Long flow = flowMap.getOrDefault(roadId, 0L);
                //累加计算车流量
                flow++;
                //将新的车流量保存到hashmap中
                flowMap.put(roadId, flow);

                //从hashmap中通过道路编号获取总的车速
                Double sumSpeed = sumSpeedMap.getOrDefault(roadId, 0.0);
                //累加计算总的车速
                sumSpeed += car.getSpeed();

                sumSpeedMap.put(roadId,sumSpeed);

                //计算平均车速
                double avgSpeed = sumSpeed / flow;
                //将结果发送到下游
                out.collect(new Result(car.getRoadId(), flow, avgSpeed));
            }
        });

        resultsDS.print();

        env.execute();

    }
}

class Result {
    private String roadId;
    private Long flow;
    private Double speed;

    @Override
    public String toString() {
        return "Result{" +
                "roadId='" + roadId + '\'' +
                ", flow=" + flow +
                ", speed=" + speed +
                '}';
    }

    public String getRoadId() {
        return roadId;
    }

    public void setRoadId(String roadId) {
        this.roadId = roadId;
    }

    public Long getFlow() {
        return flow;
    }

    public void setFlow(Long flow) {
        this.flow = flow;
    }

    public Double getSpeed() {
        return speed;
    }

    public void setSpeed(Double speed) {
        this.speed = speed;
    }

    public Result(String roadId, Long flow, Double speed) {
        this.roadId = roadId;
        this.flow = flow;
        this.speed = speed;
    }
}